import datetime
import os
from typing import List, Tuple, Dict, Any, Union
from sklearn.preprocessing import MinMaxScaler,RobustScaler
import adata
import baostock as bs
import matplotlib.dates as mdates
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
from matplotlib import gridspec
from tqdm import tqdm
from matplotlib.dates import MonthLocator

from specificStockAnalysis import SupportResistanceAnalyzer  # 导入你的分析器类
from concurrent.futures import ProcessPoolExecutor, as_completed

# 设置字体支持中文
plt.rcParams['font.sans-serif'] = ['SimHei']  # 'SimHei' 是黑体，你可以根据系统安装的字体更改
plt.rcParams['axes.unicode_minus'] = False  # 正确显示负号

# 配置参数
INITIAL_CAPITAL = 100000.0  # 初始资金
RISK_FREE_RATE = 0.03       # 无风险利率（年化）
MIN_DATA_LENGTH = 20        # 信号计算所需最小数据长度


def calculate_volatility(df, annualize=True):
    """
    计算股票的平均波动率（年化或日波动率）

    参数:
    df (pd.DataFrame): 包含收盘价的 DataFrame，列名为 'Close'
    annualize (bool): 是否年化波动率，默认为 True

    返回:
    float: 波动率（年化或日波动率）
    """
    # 计算每日收益率
    df['Daily Return'] = df['close'].pct_change()

    # 删除缺失值（第一天的收益率无法计算）
    df.dropna(inplace=True)

    # 计算日波动率（收益率的标准差）
    daily_volatility = df['Daily Return'].std()

    # 年化波动率（假设一年有 252 个交易日）
    if annualize:
        annualized_volatility = daily_volatility * np.sqrt(252)
        return annualized_volatility
    else:
        return daily_volatility


class BacktestEngine:
    def __init__(self, data: pd.DataFrame):
        self.data = data
        self.trades = []
        self.positions = []
        self.full_results = []  # 存储完整回测记录

    def run_backtest(
            self,
            risk_reward_ratio: float,
            stop_loss_pct: float = 0.02
    ) -> Tuple[Dict[str, Union[Union[float, int], Any]], List[Any], List[
        Union[Dict[str, Union[float, Any]], Dict[str, Union[float, Any]]]]]:
        """执行单次回测"""
        capital = INITIAL_CAPITAL
        portfolio = [{'date': self.data.index[0], 'value': capital}]
        self.positions.clear()
        self.trades.clear()
        for i in tqdm(range(len(self.data)), desc=f"Testing RR={risk_reward_ratio}"):
            # 获取当前数据窗口
            current_data = self.data.iloc[:i + 1]
            if len(current_data) < MIN_DATA_LENGTH:
                continue

            # 初始化信号分析器
            analyzer = SupportResistanceAnalyzer(
                data=current_data,
                bin_width=0.03
            )
            signal, vol, vol_percentile,nearest_resistance, nearest_support  = analyzer.is_buy_signal()


            # 若未持仓且触发买入信号
            if not self.positions and signal:
                entry_price = current_data['close'].iloc[-1]
                stop_price =  entry_price * (1 - stop_loss_pct)
                task_profit = entry_price * (1 + risk_reward_ratio * stop_loss_pct)
                # if nearest_resistance:
                #     # 支撑位设置为止盈点，根据stop_loss_pct盈亏比计算止损
                #     task_profit = nearest_resistance
                #     stop_price = entry_price - (task_profit-entry_price)/risk_reward_ratio
                # else:
                #     if nearest_support:
                #         # 支撑位置设为止损点，根据stop_loss_pct盈亏比计算止盈
                #         stop_price = nearest_support
                #         task_profit = entry_price + (entry_price-stop_price)*risk_reward_ratio
                position_size = capital // entry_price
                if position_size == 0:
                    continue

                # 记录持仓
                self.positions.append({
                    'entry_date': current_data.index[-1],
                    'entry_price': entry_price,
                    'stop_price': stop_price,
                    'take_profit': task_profit,
                    'size': position_size
                })

                # 记录交易
                trade = {
                    'entry_date': current_data.index[-1],
                    'entry_price': entry_price,
                    'take_profit': task_profit,
                    'stop_loss': stop_price,
                    'size': position_size,
                    'return_pct':0,
                    'status': 'open'
                }
                self.trades.append(trade)
                capital -= position_size * entry_price

            # 检查持仓退出条件
            for pos in self.positions.copy():
                current_close = current_data['close'].iloc[-1]
                if current_close <= pos['stop_price'] or current_close >= pos['take_profit']:
                    exit_price = current_close
                    pnl = (exit_price - pos['entry_price']) * pos['size']
                    capital += pos['size'] * exit_price

                    # 记录交易结果
                    self.trades.append({
                        'entry_date': pos['entry_date'],
                        'exit_date': current_data.index[-1],
                        'exit_price':exit_price,
                        'entry_price': pos['entry_price'],
                        'return_pct': (pnl / (pos['entry_price'] * pos['size'])) * 100,
                        'duration': current_data.index[-1] - pos['entry_date'],
                        'status' : 'closed'
                    })
                    self.positions.remove(pos)

            # 记录每日净值
            portfolio.append({
                'date': current_data.index[-1],
                'value': capital + sum(pos['size'] * current_data['close'].iloc[-1] for pos in self.positions)
            })

        # 计算绩效指标
        returns = pd.Series([x['value'] for x in portfolio]).pct_change().dropna()
        sharpe = (returns.mean() * 252 - RISK_FREE_RATE) / (returns.std() * np.sqrt(252))

        # ================== 新增胜率计算 ==================
        closed_trades = [t for t in self.trades if t['status'] == 'closed']
        win_rate = (sum(1 for t in closed_trades if t['return_pct'] > 0) / len(
            closed_trades)) * 100 if closed_trades else 0

        # self.plot_backtest(self.trades, portfolio)

        static_dict = {
            'risk_reward_ratio': risk_reward_ratio,
            'sharpe_ratio': sharpe,
            'total_return': (portfolio[-1]['value'] / INITIAL_CAPITAL - 1) * 100,
            'max_drawdown': self._calculate_max_drawdown([x['value'] for x in portfolio]),
            'num_trades': len(closed_trades),
            'win_rate': win_rate  # 新增胜率字段
        }


        return static_dict,self.trades,portfolio


    def _calculate_max_drawdown(self, values: List[float]) -> float:
        """计算最大回撤"""
        peak = values[0]
        max_dd = 0.0
        for v in values:
            if v > peak:
                peak = v
            dd = (peak - v) / peak
            if dd > max_dd:
                max_dd = dd
        return max_dd * 100  # 百分比形式

    def plot_backtest(self, trades: List[Dict], portfolio: List[Dict], figsize=(16, 12),save_path=None):
        closed_trades = [t for t in trades if t['status'] == 'closed']

        """可视化回测结果"""
        plt.style.use('seaborn')
        fig = plt.figure(figsize=figsize)
        gs = gridspec.GridSpec(3, 1, height_ratios=[3, 1, 2])

        # 第一子图：价格走势与交易信号
        ax1 = plt.subplot(gs[0])
        self._plot_price_signals(ax1, closed_trades)

        # 第二子图：持仓状态
        ax2 = plt.subplot(gs[1], sharex=ax1)
        self._plot_position_status(ax2, closed_trades)

        # 第三子图：资金曲线与回撤
        ax3 = plt.subplot(gs[2], sharex=ax1)
        self._plot_equity_curve(ax3, portfolio)

        plt.tight_layout()

        if save_path:
            plt.savefig(save_path)
            plt.close()
        else:
            plt.show()
    def _plot_price_signals(self, ax, trades: List[Dict]):
        # 设置日期格式化和布局
        ax.xaxis.set_major_locator(mdates.DayLocator(interval=5))  # 关键修改
        ax.xaxis.set_major_formatter(mdates.DateFormatter('%Y-%m-%d'))
        plt.xticks(rotation=45)
        plt.tight_layout()

        """绘制价格走势与交易信号"""
        # 绘制价格曲线
        ax.plot(self.data.index, self.data['close'], label='close', color='blue', linewidth=1.5)
        # --- 标注买入点并添加日期 ---
        buy_dates = [pd.to_datetime(trade['entry_date']) for trade in trades]
        buy_prices = [trade['entry_price'] for trade in trades]
        if len(buy_dates) > 0:
            # 绘制买点标记
            ax.scatter(buy_dates, buy_prices, color='limegreen', marker='^',
                       label='buy point', s=20, zorder=3)
            # 添加买点日期文本
            for date, price in zip(buy_dates, buy_prices):
                ax.text(date, price,
                        date.strftime("%m-%d"),  # 显示月-日格式
                        color='darkgreen',
                        fontsize=8,
                        ha='left',  # 水平对齐：左对齐
                        va='bottom',  # 垂直对齐：底部对齐
                        rotation=0,  # 轻微旋转防重叠
                        bbox=dict(facecolor='white', alpha=0.7, edgecolor='none'))  # 白色背景框

        # --- 标注卖出点并添加日期 ---
        sell_dates = [pd.to_datetime(trade['exit_date']) for trade in trades]
        sell_prices = [trade['exit_price'] for trade in trades]
        if len(sell_dates) > 0:
            # 绘制卖点标记
            ax.scatter(sell_dates, sell_prices, color='red', marker='v',
                       label='sell point', s=20, zorder=3)
            # 添加卖点日期文本
            for date, price in zip(sell_dates, sell_prices):
                ax.text(date, price,
                        date.strftime("%m-%d"),
                        color='darkred',
                        fontsize=8,
                        ha='right',  # 水平对齐：右对齐
                        va='top',  # 垂直对齐：顶部对齐
                        rotation=0,  # 反向旋转
                        bbox=dict(facecolor='white', alpha=0.7, edgecolor='none'))


        ax.set_title('price_trend_and_trading_signal')
        ax.set_ylabel('price')
        ax.legend(loc='upper left')

    def _plot_position_status(self,ax, trades: List[Dict]):
        """绘制持仓状态（盈利红色/亏损绿色）"""
        # 设置坐标轴格式
        ax.xaxis.set_major_locator(mdates.DayLocator(interval=5))  # 修正括号
        ax.xaxis.set_major_formatter(mdates.DateFormatter('%Y-%m-%d'))
        plt.xticks(rotation=45)
        plt.tight_layout()

        # 生成持仓状态序列
        position = pd.Series(0, index=self.data.index)

        # 遍历每个交易单独绘制颜色
        for trade in trades:
            # 获取交易时间段
            entry_date = pd.to_datetime(trade['entry_date'])
            exit_date = pd.to_datetime(trade['exit_date'])

            # 计算盈亏（假设有exit_price和entry_price字段）
            profit = trade['exit_price'] - trade['entry_price']
            color = 'red' if profit > 0 else 'green'  # 盈利红/亏损绿

            # 标记持仓区间
            mask = (self.data.index >= entry_date) & (self.data.index <= exit_date)

            # 绘制单个交易的持仓区间
            ax.fill_between(self.data.index[mask], 1, 0,
                            where=(position[mask] >= 0),
                            facecolor=color,
                            alpha=0.3,
                            edgecolor='none')

        ax.set_ylabel('Position Status')
        ax.set_yticks([0, 1])
        ax.set_yticklabels(['No Position', 'In Position'])

    def _plot_equity_curve(self, ax, portfolio: List[Dict]):
        ax.xaxis.set_major_locator(mdates.DayLocator(interval=5))  # 关键修改
        ax.xaxis.set_major_formatter(mdates.DateFormatter('%Y-%m-%d'))
        plt.xticks(rotation=45)
        plt.tight_layout()
        """绘制资金曲线与回撤"""
        # 资金曲线
        portfolio_df = pd.DataFrame(portfolio)
        portfolio_df['date'] = pd.to_datetime(portfolio_df['date'])
        ax.plot(portfolio_df['date'], portfolio_df['value'],
                label='net value', color='blue', linewidth=1.5)

        # 计算回撤
        peak = portfolio_df['value'].cummax()
        drawdown = (peak - portfolio_df['value']) / peak
        ax.fill_between(portfolio_df['date'], drawdown,
                        color='red', alpha=0.3, label='draw down')

        ax.set_ylabel('account net value')
        ax.legend(loc='upper left')
def get_past_date(n_days):
    current_time = datetime.datetime.now()
    past_date = current_time - datetime.timedelta(days=n_days)
    return past_date.strftime('%Y-%m-%d')

def load_data(data) -> pd.DataFrame:
    """加载股票数据，确保格式正确"""
    # trade_date 列转成datetime格式并设置为索引
    data['trade_date'] = pd.to_datetime(data['trade_date'])
    data.set_index('trade_date', inplace=True)
    return data[['open', 'high', 'low', 'close', 'volume']]
def optimize_parameters(data, rr_range: List[float], stop_loss_pct: float = 0.02,save_path = None) -> pd.DataFrame:
    """多参数优化"""
    data = load_data(data)
    engine = BacktestEngine(data)
    results = []
    trades = []
    portforios = []

    for rr in rr_range:
        result, trade, portfolio = engine.run_backtest(risk_reward_ratio=rr,stop_loss_pct = stop_loss_pct)
        results.append(result)
        trades.append(trade)
        portforios.append(portfolio)

    # 选出夏普比率最大的result对应的index， 并绘制回测结果图
    best_index = np.argmax([result['sharpe_ratio'] for result in results])
    if save_path:
        engine.plot_backtest(trades[best_index], portforios[best_index], save_path=save_path)
    else:
        engine.plot_backtest(trades[best_index], portforios[best_index])


    df = pd.DataFrame(results)
    # df.to_csv('optimization_results_.csv', index=False)
    return df





def select_top_stocks(csv_path, top_n=10):
    """
    多维度选股算法
    :param csv_path: 数据文件路径
    :param top_n: 选取前N名
    :return: 优选股票DataFrame
    """
    # 数据读取与预处理
    df = pd.read_csv(csv_path)
    required_columns = ['sharpe_ratio', 'win_rate', 'num_trades', 'total_return']
    df = df.dropna(subset=required_columns)  # 去除缺失值
    # 转换为字符串类型
    df['stock_code'] = df['stock_code'].astype(str)

    # 使用zfill补零
    df['stock_code'] = df['stock_code'].str.zfill(6)

    # 数据标准化
    scaler = RobustScaler()
    normalized = scaler.fit_transform(df[required_columns])
    normalized_df = pd.DataFrame(normalized, columns=required_columns)

    # 动态权重配置（可根据需求调整）
    weights = {
        'sharpe_ratio': 0.35,  # 夏普比率最重要
        'win_rate': 0.25,  # 胜率次重要
        'total_return': 0.25,  # 总回报同样重要
        'num_trades': 0.15  # 交易次数参考权重
    }

    # 计算综合评分
    normalized_df['composite_score'] = (
            normalized_df['sharpe_ratio'] * weights['sharpe_ratio'] +
            normalized_df['win_rate'] * weights['win_rate'] +
            normalized_df['total_return'] * weights['total_return'] +
            normalized_df['num_trades'] * weights['num_trades']
    )

    # 合并原始数据
    result_df = df[['stock_code', 'stock_name','sharpe_ratio','win_rate','total_return']].copy()
    result_df['composite_score'] = normalized_df['composite_score']

    # 分层筛选：确保各维度基础表现
    filtered = result_df[
        (normalized_df['sharpe_ratio'] > 0.4) &  # 夏普比率前40%
        (normalized_df['win_rate'] > 0.5) &  # 胜率超过中位数
        (normalized_df['total_return'] > 0.3)&  # 总回报前30%
        (normalized_df['num_trades'] > 0.5) # 交易次数前50%
        ]

    # 排序取前N名
    top_stocks = filtered.sort_values('composite_score', ascending=False).head(top_n)

    return top_stocks.reset_index(drop=True)


def run_backtest_with_params(stock_code, start_date, rr_range,save_path = None):
    res_df = adata.stock.market.get_market(stock_code=stock_code, k_type=1, start_date=start_date)
    yearly_return = calculate_volatility(res_df, False)


    """辅助函数，用于多线程调用"""

    result = optimize_parameters(res_df, rr_range, stop_loss_pct=yearly_return,save_path = save_path)
    return result


if __name__ == "__main__":
    # 示例参数范围
    rr_range = np.arange(0.8, 2.0, 0.1).tolist()

    adata.proxy(is_proxy=True, ip='127.0.0.1:10810')
    # 示例：获取3天前的日期
    n = 400
    top_n = 500 # 监控市场上的500只票
    start_date = get_past_date(n)


    res_df = adata.stock.info.all_code()
    # 组合 exchange 和 stock_code 成 "SZ.00001" 形式的列表
    # 创建字典，将exchange和stock_code组合成key，short_name作为value
    combined_map = {f"{row['exchange']}.{row['stock_code']}": row['short_name'] for index, row in res_df.iterrows()}
    stock_code_name_map = {f"{row['stock_code']}": row['short_name'] for index, row in res_df.iterrows()}


    colume_list = ['stock_code','stock_name','risk_reward_ratio','sharpe_ratio','total_return','max_drawdown','num_trades','win_rate','has_signal_today','volatility']
    total_df = pd.DataFrame(columns=colume_list)


    #传入一个给定的csv，包含股票的 stock_code, stock_name, risk_reward_ratio, sharpe_ratio, total_return, max_drawdown, num_trades, win_rate, volatility 等统计信息，要求根据
    csv_path = 'strategy_first.csv'
    df = select_top_stocks(csv_path, top_n)

    candidate_list = []
    base_dir = "./"
    # 遍历每个股票，找出当日出现买入信号的股票，并保存到列表中
    for stock_code in df['stock_code']:
        res_df = adata.stock.market.get_market(stock_code=stock_code, k_type=1, start_date=start_date)
        if res_df.empty:
            continue
        yearly_return = calculate_volatility(res_df, False)
        stock_name = stock_code_name_map[stock_code]
        save_path_base = os.path.join(base_dir, "./pictures/{} 支撑压力分析".format(stock_name))
        save_path_radar = os.path.join(base_dir, "./pictures/{} 信号雷达图".format(stock_name))
        back_test_dir = os.path.join(base_dir, "./pictures/{} 回测结果".format(stock_name))

        # 示例1：使用在线数据
        analyzer = SupportResistanceAnalyzer(
            data=res_df,
            bin_width=0.03
        )
        # 获取信号数据
        signal, strength_ratio, distance_ratio, _, _ = analyzer.is_buy_signal()
        if signal:
            print(f"{stock_code} {stock_code_name_map[stock_code]} 触发买入信号")
            candidate_list.append(stock_code)
            # 绘制支撑压力位置
            analyzer.plot(save_path=save_path_base)
            # 绘制信号雷达图
            analyzer.plot_signal_strength(save_path=save_path_radar)
            # 优化该股票的参数，绘制股票的支撑压力位置，以及信号雷达图，以及回测结果，保存到pictures文件夹中
            result = run_backtest_with_params(stock_code, start_date, rr_range,save_path=back_test_dir)
            if result is not None:
                print(f"{stock_code} {stock_code_name_map[stock_code]} 优化参数结果：")
                # 输出最佳参数
                best_row = result.loc[result['sharpe_ratio'].idxmax()]
                best_row['stock_code'] = stock_code
                best_row['stock_name'] = stock_code_name_map[stock_code]
                best_row['volatility'] = yearly_return
                total_df = pd.concat([total_df, best_row.to_frame().T], ignore_index=True)

                # 更新后的输出格式
                print(f"""
                ================ {stock_code}优化结果 ================
                最佳盈亏比:      {best_row['risk_reward_ratio']:.2f}
                夏普比率:       {best_row['sharpe_ratio']:.2f}
                总收益率:       {best_row['total_return']:.2f}%
                最大回撤:       {best_row['max_drawdown']:.2f}%
                胜率:          {best_row['win_rate']:.1f}%
                交易次数:       {best_row['num_trades']}
                日波动率：       {yearly_return:.2f}%
                ========================================
                """)

    # stock_code = "688097"
    # res_df = adata.stock.market.get_market(stock_code=stock_code, k_type=1, start_date=start_date)
    #
    # yearly_return = calculate_volatility(res_df, False)
    # #运行优化
    # results_df = optimize_parameters(res_df, rr_range, yearly_return)
    #
    # 输出最佳参数
    # best_row = results_df.loc[results_df['sharpe_ratio'].idxmax()]
    # best_row['stock_code'] = stock_code
    # best_row['stock_name'] = stock_code_name_map[stock_code]
    # best_row['volatility'] = yearly_return
    # total_df = pd.concat([total_df, best_row.to_frame().T], ignore_index=True)
    #
    # # 更新后的输出格式
    # print(f"""
    # ================ {stock_code}优化结果 ================
    # 最佳盈亏比:      {best_row['risk_reward_ratio']:.2f}
    # 夏普比率:       {best_row['sharpe_ratio']:.2f}
    # 总收益率:       {best_row['total_return']:.2f}%
    # 最大回撤:       {best_row['max_drawdown']:.2f}%
    # 胜率:          {best_row['win_rate']:.1f}%
    # 交易次数:       {best_row['num_trades']}
    # 日波动率：       {yearly_return:.2f}%
    # ========================================
    # """)

